package whut.pso;

import org.apache.commons.math3.linear.ArrayRealVector;
import org.apache.commons.math3.linear.RealVector;
import whut.commons.BaseBlock;
import whut.commons.IProblemSet;
import whut.commons.Location;

import java.util.Arrays;
import java.util.Random;
import java.util.Vector;

/**
 * Created by Justin on 2018-04-02
 */
public class PSOProcess extends BaseBlock implements IPSOConstants, IProcess{
    protected Vector<Particle> swarm;
    protected double[] pBest;
    protected Vector<Location> pBestLocation;
    protected double gBest;
    protected Location gBestLocation;
    protected double[] fitnessValueList;

    protected Random generator;

    protected IProblemSet problem;
    protected INormalization normalize;
    protected IWeightHandler weightHandler;
    protected IUpdateSwarm updateSwarm;

    public PSOProcess(IProblemSet problem) {
        this(problem,null, null,null);
    }

    public PSOProcess(IProblemSet problem, INormalization normalize) {
        this(problem, normalize, null,null);
    }

    public PSOProcess(IProblemSet problem, INormalization normalize, IWeightHandler weightHandler) {
        this(problem, normalize, weightHandler, null);
    }

    public PSOProcess(IProblemSet problem, INormalization normalize, IWeightHandler weightHandler,
                      IUpdateSwarm updateSwarm) {
        this.problem = problem;
        this.normalize = normalize;
        this.weightHandler = weightHandler;
        this.updateSwarm = updateSwarm;
        generator = new Random();
        pBest = new double[SWARM_SIZE];
        swarm = new Vector<Particle>();
        pBestLocation = new Vector<Location>();
        fitnessValueList = new double[SWARM_SIZE];
        gBest = Double.MAX_VALUE;

        gBestLocation = null;
    }

    public void clear(){
        Arrays.fill(pBest, 0);
        Arrays.fill(fitnessValueList, 0);
        swarm.clear();
        pBestLocation.clear();
        gBest = Double.MAX_VALUE;
        gBestLocation = null;
    }
    
    @Override
    public Location execute() {
        clear();

        initializeSwarm(null, 0);
        updateFitnessList();

        for (int i = 0; i < SWARM_SIZE; i++) {
            pBest[i] = fitnessValueList[i];
            pBestLocation.add(swarm.get(i).getLocation());
        }

        return searching();
    }

    @Override
    public Location searching() {
        int t = 0;
        double w;
        double err = 9999;

        while (t < MAX_ITERATION && err > problem.getERR_TOLERANCE()) {
            // step 1 - update pBest
            for (int i = 0; i < SWARM_SIZE; i++) {
                if (fitnessValueList[i] < pBest[i]) {
                    pBest[i] = fitnessValueList[i];
                    pBestLocation.set(i, swarm.get(i).getLocation());
                }
            }

            // step 2 - update gBest
            int bestParticleIndex = PSOUtility.getMinPos(fitnessValueList);
            if (t == 0 || fitnessValueList[bestParticleIndex] < gBest) {
                gBest = fitnessValueList[bestParticleIndex];
                gBestLocation = swarm.get(bestParticleIndex).getLocation();
            }
            if (weightHandler == null){
                w = W_UPPERBOUND - (((double) t) / MAX_ITERATION) * (W_UPPERBOUND - W_LOWERBOUND);
            }else{
                w = weightHandler.getWeight(t);
            }

            for (int i = 0; i < SWARM_SIZE; i++) {

                Particle p = swarm.get(i);
                double[] newLoc;

                if (updateSwarm == null){
                    // step 3 - update velocity
                    double[] newVel = updateVelocity(p, pBestLocation.get(i).getLoc(),
                            gBestLocation.getLoc(), w);
                    Velocity vel = new Velocity(newVel);
                    p.setVelocity(vel);

                    // step 4 - update location
                    newLoc = updateLocation(p, newVel);
                }else{
                    // step 4 - update location
                    newLoc = updateSwarm.updateSwarm(p.getLocation().getLoc(),  pBestLocation.get(i).getLoc(),
                            gBestLocation.getLoc(), w, normalize, problem);
                }
                // step 5 - after update location,方便扩展
                newLoc = afterUpdateLocation(p, newLoc);
                Location loc = new Location(newLoc);
                p.setLocation(loc);
            }

            err = problem.evaluate(gBestLocation) - 0; // minimizing the functions means it's getting closer to 0

            mLog.trace("ITERATION {} :\t Err = {} \tBest: {}", t, err, gBestLocation.toString() );
            t++;
            updateFitnessList();
        }

        if (err > problem.getERR_TOLERANCE()){
            mLog.error("Solution didn't found in {} iterations, the solutions is:\t Best:  {}",
                    t , gBestLocation.toString());
            return null;
        }else{
            mLog.info("Solution found at iteration {}, the solutions is: \t Best:  {}",
                    t - 1, gBestLocation.toString());
            return gBestLocation;
        }
    }

    protected double[] afterUpdateLocation(Particle p, double[] value){
        return value;
    }

    @Override
    public void initializeSwarm(double[] center, double range) {
        Particle p;
        for (int i = 0; i < SWARM_SIZE; i++) {
            p = new Particle();

            int PROBLEM_DIMENSION = problem.getPROBLEM_DIMENSION();
            // randomize location inside a space defined in Problem Set
            double[] loc = new double[PROBLEM_DIMENSION];
            for (int j = 0; j < PROBLEM_DIMENSION; j++) {
                if (center == null){
                    loc[j] = problem.getLOC_LOW() + generator.nextDouble() * (problem.getLOC_HIGH() - problem.getLOC_LOW());
                }else{
                    loc[j] = center[j] + generator.nextDouble() * range;
                }
            }
            if (normalize != null){
                loc = normalize.normalizeLocationOfRange(loc, center, range);
            }
            Location location = new Location(loc);

            // randomize velocity in the range defined in Problem Set
            double[] vel = new double[PROBLEM_DIMENSION];
            for (int j = 0; j < PROBLEM_DIMENSION; j++) {
                vel[j] = problem.getVEL_LOW() + generator.nextDouble() * (problem.getVEL_HIGH() - problem.getVEL_LOW());
            }
            Velocity velocity = new Velocity(vel);

            p.setLocation(location);
            p.setVelocity(velocity);
            swarm.add(p);
        }
    }

    @Override
    public void updateFitnessList() {
        for (int i = 0; i < SWARM_SIZE; i++) {
            fitnessValueList[i] = swarm.get(i).getFitnessValue(problem);
        }
    }

    @Override
    public double[] updateVelocity(Particle p, double[] pBest, double[] gBest, double w) {
        RealVector pL = new ArrayRealVector(p.getLocation().getLoc());
        RealVector pBestLoc = new ArrayRealVector(pBest);
        RealVector gBestLoc = new ArrayRealVector(gBest);

        double r1 = generator.nextDouble();
        double r2 = generator.nextDouble();

        RealVector pV = new ArrayRealVector(p.getVelocity().getPos());
        RealVector diffPBest = pBestLoc.subtract(pL).mapMultiply(r1 * C1);
        RealVector diffGBeat = gBestLoc.subtract(pL).mapMultiply(r2 * C2);
        RealVector result = (pV.mapMultiply(w)).add(diffPBest).add(diffGBeat);

        return result.toArray();
    }

    @Override
    public double[] updateLocation(Particle p, double[] v) {
        RealVector pL = new ArrayRealVector(p.getLocation().getLoc());
        RealVector newLoc =  pL.add(new ArrayRealVector(v));
        
        if (normalize != null){
            return normalize.normalizeLocationOfRange(newLoc.toArray(), null, 0);
        }
        return newLoc.toArray();
    }
}
